#!/usr/bin/env python3
"""
地理勘查记录系统性能测试脚本
"""

import requests
import time
import threading
import statistics
import json
from datetime import datetime, timedelta
from concurrent.futures import ThreadPoolExecutor, as_completed

class PerformanceTester:
    def __init__(self, base_url='http://127.0.0.1:8000'):
        self.base_url = base_url
        self.access_token = None
        self.results = []
        
    def setup_auth(self):
        """设置认证token"""
        # 创建测试用户
        register_url = f"{self.base_url}/api/users/register/"
        user_data = {
            "username": f"perftest_{int(time.time())}",
            "email": f"perftest_{int(time.time())}@example.com",
            "phone": "13888888888",
            "password": "testpass123",
            "password_confirm": "testpass123"
        }
        
        try:
            response = requests.post(register_url, json=user_data)
            if response.status_code == 201:
                result = response.json()
                self.access_token = result.get('access_token')
                return True
        except:
            pass
        
        # 如果注册失败，尝试登录
        login_url = f"{self.base_url}/api/auth/token/"
        login_data = {
            "username": "testuser",
            "password": "testpass123"
        }
        
        try:
            response = requests.post(login_url, json=login_data)
            if response.status_code == 200:
                result = response.json()
                self.access_token = result.get('access')
                return True
        except:
            pass
        
        return False
    
    def get_headers(self):
        """获取请求头"""
        return {
            'Authorization': f'Bearer {self.access_token}',
            'Content-Type': 'application/json'
        }
    
    def create_test_data(self, count=100):
        """创建测试数据"""
        print(f"📝 创建 {count} 条测试数据...")
        
        url = f"{self.base_url}/api/surveys/"
        headers = self.get_headers()
        
        created = 0
        for i in range(count):
            data = {
                "location_name": f"性能测试地点 {i+1}",
                "address": f"测试地址 {i+1}",
                "latitude": str(39.908823 + (i * 0.001)),
                "longitude": str(116.397470 + (i * 0.001)),
                "orientation": "坐北朝南",
                "notes": f"性能测试备注 {i+1}",
                "survey_date": (datetime.now() - timedelta(days=i % 30)).isoformat()
            }
            
            try:
                response = requests.post(url, json=data, headers=headers, timeout=5)
                if response.status_code == 201:
                    created += 1
            except:
                pass
            
            if (i + 1) % 20 == 0:
                print(f"   已创建 {created}/{i+1} 条记录")
        
        print(f"✅ 成功创建 {created} 条测试数据")
        return created
    
    def measure_response_time(self, url, method='GET', data=None, headers=None):
        """测量单个请求的响应时间"""
        start_time = time.time()
        
        try:
            if method == 'GET':
                response = requests.get(url, headers=headers, timeout=10)
            elif method == 'POST':
                response = requests.post(url, json=data, headers=headers, timeout=10)
            else:
                return None, None
            
            end_time = time.time()
            response_time = (end_time - start_time) * 1000  # 转换为毫秒
            
            return response_time, response.status_code
        except Exception as e:
            return None, str(e)
    
    def test_endpoint_performance(self, endpoint, method='GET', data=None, requests_count=50, concurrent=10):
        """测试单个端点的性能"""
        print(f"\n🔧 测试端点: {method} {endpoint}")
        print(f"   请求数量: {requests_count}, 并发数: {concurrent}")
        
        url = f"{self.base_url}{endpoint}"
        headers = self.get_headers()
        
        response_times = []
        status_codes = []
        errors = []
        
        def single_request():
            response_time, status = self.measure_response_time(url, method, data, headers)
            if response_time is not None:
                response_times.append(response_time)
                status_codes.append(status)
            else:
                errors.append(status)
        
        # 并发测试
        start_time = time.time()
        
        with ThreadPoolExecutor(max_workers=concurrent) as executor:
            futures = [executor.submit(single_request) for _ in range(requests_count)]
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    errors.append(str(e))
        
        total_time = time.time() - start_time
        
        # 计算统计数据
        if response_times:
            avg_time = statistics.mean(response_times)
            median_time = statistics.median(response_times)
            min_time = min(response_times)
            max_time = max(response_times)
            p95_time = statistics.quantiles(response_times, n=20)[18] if len(response_times) >= 20 else max_time
            
            success_rate = len([s for s in status_codes if 200 <= s < 300]) / len(status_codes) * 100
            qps = requests_count / total_time
            
            # 打印结果
            print(f"   📊 性能指标:")
            print(f"      平均响应时间: {avg_time:.2f} ms")
            print(f"      中位数响应时间: {median_time:.2f} ms")
            print(f"      最小响应时间: {min_time:.2f} ms")
            print(f"      最大响应时间: {max_time:.2f} ms")
            print(f"      95% 响应时间: {p95_time:.2f} ms")
            print(f"      成功率: {success_rate:.1f}%")
            print(f"      QPS: {qps:.2f}")
            
            if errors:
                print(f"      错误数量: {len(errors)}")
            
            return {
                'endpoint': endpoint,
                'method': method,
                'requests_count': requests_count,
                'concurrent': concurrent,
                'avg_time': avg_time,
                'median_time': median_time,
                'min_time': min_time,
                'max_time': max_time,
                'p95_time': p95_time,
                'success_rate': success_rate,
                'qps': qps,
                'errors': len(errors)
            }
        else:
            print(f"   ❌ 所有请求都失败了")
            return None
    
    def test_list_performance(self):
        """测试列表查询性能"""
        return self.test_endpoint_performance('/api/surveys/', 'GET', requests_count=100, concurrent=20)
    
    def test_search_performance(self):
        """测试搜索性能"""
        endpoint = '/api/surveys/search/?search=测试&start_date=2024-01-01&end_date=2024-12-31'
        return self.test_endpoint_performance(endpoint, 'GET', requests_count=50, concurrent=10)
    
    def test_statistics_performance(self):
        """测试统计接口性能"""
        return self.test_endpoint_performance('/api/surveys/statistics/', 'GET', requests_count=30, concurrent=5)
    
    def test_create_performance(self):
        """测试创建记录性能"""
        data = {
            "location_name": "性能测试地点",
            "address": "性能测试地址",
            "latitude": "39.908823",
            "longitude": "116.397470",
            "orientation": "坐北朝南",
            "notes": "性能测试备注",
            "survey_date": datetime.now().isoformat()
        }
        
        return self.test_endpoint_performance('/api/surveys/', 'POST', data, requests_count=50, concurrent=5)
    
    def test_concurrent_users(self, user_count=10, requests_per_user=20):
        """测试并发用户场景"""
        print(f"\n🔥 并发用户测试: {user_count} 个用户，每用户 {requests_per_user} 请求")
        
        def user_session():
            """模拟单个用户会话"""
            session_times = []
            session = requests.Session()
            session.headers.update(self.get_headers())
            
            for _ in range(requests_per_user):
                # 模拟用户行为：查看列表 -> 搜索 -> 查看详情
                endpoints = [
                    f"{self.base_url}/api/surveys/",
                    f"{self.base_url}/api/surveys/search/?search=测试",
                    f"{self.base_url}/api/surveys/statistics/"
                ]
                
                for endpoint in endpoints:
                    start_time = time.time()
                    try:
                        response = session.get(endpoint, timeout=10)
                        response_time = (time.time() - start_time) * 1000
                        session_times.append(response_time)
                        
                        # 随机延迟，模拟真实用户行为
                        time.sleep(0.1)
                    except:
                        pass
            
            return session_times
        
        # 并发执行用户会话
        all_times = []
        start_time = time.time()
        
        with ThreadPoolExecutor(max_workers=user_count) as executor:
            futures = [executor.submit(user_session) for _ in range(user_count)]
            for future in as_completed(futures):
                try:
                    times = future.result()
                    all_times.extend(times)
                except Exception as e:
                    print(f"   用户会话失败: {e}")
        
        total_time = time.time() - start_time
        
        if all_times:
            avg_time = statistics.mean(all_times)
            total_requests = len(all_times)
            overall_qps = total_requests / total_time
            
            print(f"   📊 并发测试结果:")
            print(f"      总请求数: {total_requests}")
            print(f"      平均响应时间: {avg_time:.2f} ms")
            print(f"      总耗时: {total_time:.2f} s")
            print(f"      整体QPS: {overall_qps:.2f}")
            
            return {
                'user_count': user_count,
                'requests_per_user': requests_per_user,
                'total_requests': total_requests,
                'avg_time': avg_time,
                'total_time': total_time,
                'qps': overall_qps
            }
        
        return None
    
    def run_performance_tests(self):
        """运行所有性能测试"""
        print("🚀 开始地理勘查记录系统性能测试")
        print("=" * 60)
        
        # 设置认证
        if not self.setup_auth():
            print("❌ 认证设置失败")
            return False
        
        print("✅ 认证设置成功")
        
        # 创建测试数据
        self.create_test_data(200)
        
        # 性能测试项目
        tests = [
            ("列表查询性能", self.test_list_performance),
            ("搜索查询性能", self.test_search_performance),
            ("统计接口性能", self.test_statistics_performance),
            ("创建记录性能", self.test_create_performance),
        ]
        
        results = []
        
        for test_name, test_func in tests:
            print(f"\n🔧 执行性能测试: {test_name}")
            try:
                result = test_func()
                if result:
                    results.append(result)
            except Exception as e:
                print(f"   ❌ 测试失败: {e}")
        
        # 并发用户测试
        concurrent_result = self.test_concurrent_users(10, 10)
        if concurrent_result:
            results.append(concurrent_result)
        
        # 生成性能报告
        self.generate_performance_report(results)
        
        return True
    
    def generate_performance_report(self, results):
        """生成性能测试报告"""
        print("\n" + "=" * 60)
        print("📈 性能测试报告")
        print("=" * 60)
        
        # 性能基准
        benchmarks = {
            'avg_time': 500,  # 平均响应时间 < 500ms
            'p95_time': 1000,  # 95%响应时间 < 1000ms
            'success_rate': 99.0,  # 成功率 > 99%
            'qps': 10.0  # QPS > 10
        }
        
        passed_tests = 0
        total_tests = 0
        
        for result in results:
            if 'endpoint' in result:
                print(f"\n📊 {result['endpoint']} ({result['method']})")
                print(f"   平均响应时间: {result['avg_time']:.2f} ms {'✅' if result['avg_time'] < benchmarks['avg_time'] else '❌'}")
                print(f"   95%响应时间: {result['p95_time']:.2f} ms {'✅' if result['p95_time'] < benchmarks['p95_time'] else '❌'}")
                print(f"   成功率: {result['success_rate']:.1f}% {'✅' if result['success_rate'] >= benchmarks['success_rate'] else '❌'}")
                print(f"   QPS: {result['qps']:.2f} {'✅' if result['qps'] >= benchmarks['qps'] else '❌'}")
                
                # 计算通过的指标数量
                passed = sum([
                    result['avg_time'] < benchmarks['avg_time'],
                    result['p95_time'] < benchmarks['p95_time'],
                    result['success_rate'] >= benchmarks['success_rate'],
                    result['qps'] >= benchmarks['qps']
                ])
                
                if passed >= 3:  # 至少3个指标通过
                    passed_tests += 1
                total_tests += 1
        
        print(f"\n🏆 性能测试总结:")
        print(f"   通过测试: {passed_tests}/{total_tests}")
        
        if passed_tests == total_tests:
            print("   🎉 所有性能测试通过！系统性能良好")
        elif passed_tests >= total_tests * 0.7:
            print("   ⚠️  大部分测试通过，系统性能可接受")
        else:
            print("   ❌ 多个测试未通过，需要优化系统性能")
        
        # 保存详细报告
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        report_file = f"performance_report_{timestamp}.json"
        
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        
        print(f"   📄 详细报告已保存至: {report_file}")


def main():
    """主函数"""
    import sys
    
    base_url = sys.argv[1] if len(sys.argv) > 1 else 'http://127.0.0.1:8000'
    
    print(f"🔧 目标服务器: {base_url}")
    
    tester = PerformanceTester(base_url)
    success = tester.run_performance_tests()
    
    sys.exit(0 if success else 1)


if __name__ == '__main__':
    main()